{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "06e5af38-85ad-46b6-a5ea-52a7ff947853",
   "metadata": {},
   "source": [
    "# Labelme转YOLO-单个文件\n",
    "\n",
    "同济子豪兄 2023-4-16"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0a8a2a1-fb46-4d51-8b4f-cd0568bc13f8",
   "metadata": {},
   "source": [
    "## 导入工具包"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5969d148-19e6-4b2e-9797-01105b87b56c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31ef8f72-7733-4d84-8250-c48583f416a3",
   "metadata": {},
   "source": [
    "## 数据集类别信息"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e3125fd2-58b3-42b9-a1b5-726a195f923e",
   "metadata": {},
   "outputs": [],
   "source": [
    "classes = {\n",
    "    'Angle':0,\n",
    "    'Ruler':1,\n",
    "    'Triangle_30':2,\n",
    "    'Triangle_45':3\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f2c5576c-f76f-4d37-8101-08ea62856e68",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Angle', 'Ruler', 'Triangle_30', 'Triangle_45']"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(classes.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8aa499c5-f1ae-4f81-b5f9-a64ff727986e",
   "metadata": {},
   "source": [
    "## 生成`classes.txt`文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "fff6141f-70af-4fe2-9e4c-618edeefd7e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('classes.txt', 'w', encoding='utf-8') as f:\n",
    "    for each in list(classes.keys()):\n",
    "        f.write(each + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b4fb7ee-0886-416c-bbaf-6c2dfe7e63c9",
   "metadata": {},
   "source": [
    "## 载入一个labelme格式的json标注文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "475e8269-be3e-48bf-b089-9538f63c4b19",
   "metadata": {},
   "outputs": [],
   "source": [
    "labelme_path = '1.json'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b17d0e7a-b987-4784-9dfe-a67ebfeefb15",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(labelme_path, 'r', encoding='utf-8') as f:\n",
    "    labelme = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "2c2b711c-4e34-459b-854f-bfdda7170935",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['version', 'flags', 'shapes', 'imagePath', 'imageData', 'imageHeight', 'imageWidth'])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labelme.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "979f8578-9c1c-4b32-b98e-53838d0ef5fe",
   "metadata": {},
   "source": [
    "### 元数据 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9fb38995-4ecc-44b4-a132-1d6851f18961",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'5.1.1'"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labelme['version']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "376d7a49-d5a1-4827-b68e-7334a231a144",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{}"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labelme['flags']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "1733f8f9-72d4-47f8-a010-2ab98b724afa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.jpg'"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 图像文件名\n",
    "labelme['imagePath']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "c0411a3c-8e43-4878-8657-cbdca9e15c18",
   "metadata": {},
   "outputs": [],
   "source": [
    "labelme['imageData']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "9bc00c97-9d5f-4664-b502-631cdf2a15b8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2736"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 图像高度\n",
    "labelme['imageHeight']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "8185cf9e-84ba-4081-8e4a-db189b10d29e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3648"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 图像宽度\n",
    "labelme['imageWidth']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68bd2fd7-914f-4ebc-a65e-bf43ceb87761",
   "metadata": {},
   "source": [
    "### 标注信息"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "2ddb2f62-b076-491f-83cf-aa4f12c82ea3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'label': 'Triangle_30',\n",
       "  'points': [[172.52941176470614, 100.35294117647076],\n",
       "   [2054.8823529411766, 976.8235294117648]],\n",
       "  'group_id': None,\n",
       "  'shape_type': 'rectangle',\n",
       "  'flags': {}},\n",
       " {'label': 'Triangle_45',\n",
       "  'points': [[2422.529411764706, 432.70588235294133],\n",
       "   [3569.5882352941176, 1673.8823529411766]],\n",
       "  'group_id': None,\n",
       "  'shape_type': 'rectangle',\n",
       "  'flags': {}},\n",
       " {'label': 'Angle',\n",
       "  'points': [[1128.4117647058827, 1000.3529411764706],\n",
       "   [2510.764705882353, 1844.4705882352944]],\n",
       "  'group_id': None,\n",
       "  'shape_type': 'rectangle',\n",
       "  'flags': {}},\n",
       " {'label': 'Ruler',\n",
       "  'points': [[290.17647058823553, 2029.764705882353],\n",
       "   [2510.764705882353, 2626.8235294117644]],\n",
       "  'group_id': None,\n",
       "  'shape_type': 'rectangle',\n",
       "  'flags': {}}]"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labelme['shapes']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afca9bb8-8197-43d2-82b6-9a7e2b6f8dc2",
   "metadata": {},
   "source": [
    "## 生成YOLO格式的标注文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "eaf81923-f9c2-4803-a285-f838ef672729",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.json --> 1.txt 转换完成\n"
     ]
    }
   ],
   "source": [
    "img_width = labelme['imageWidth']   # 图像宽度\n",
    "img_height = labelme['imageHeight'] # 图像高度\n",
    "\n",
    "# 生成 YOLO 格式的 txt 文件\n",
    "suffix = labelme_path.split('.')[-2]\n",
    "yolo_txt_path = suffix + '.txt'\n",
    "\n",
    "with open(yolo_txt_path, 'w', encoding='utf-8') as f:\n",
    "    for each_ann in labelme['shapes']: # 遍历每个框\n",
    "\n",
    "        if each_ann['shape_type'] == 'rectangle': # 筛选出框\n",
    "\n",
    "            # 获取类别 ID\n",
    "            bbox_class_id = classes[each_ann['label']]\n",
    "\n",
    "            # 左上角和右下角的 XY 像素坐标\n",
    "            bbox_top_left_x = int(min(each_ann['points'][0][0], each_ann['points'][1][0]))\n",
    "            bbox_bottom_right_x = int(max(each_ann['points'][0][0], each_ann['points'][1][0]))\n",
    "            bbox_top_left_y = int(min(each_ann['points'][0][1], each_ann['points'][1][1]))\n",
    "            bbox_bottom_right_y = int(max(each_ann['points'][0][1], each_ann['points'][1][1]))\n",
    "\n",
    "            # 框中心点的 XY 像素坐标\n",
    "            bbox_center_x = int((bbox_top_left_x + bbox_bottom_right_x) / 2)\n",
    "            bbox_center_y = int((bbox_top_left_y + bbox_bottom_right_y) / 2)\n",
    "\n",
    "            # 框宽度\n",
    "            bbox_width = bbox_bottom_right_x - bbox_top_left_x\n",
    "\n",
    "            # 框高度\n",
    "            bbox_height = bbox_bottom_right_y - bbox_top_left_y\n",
    "\n",
    "            # 框中心点归一化坐标\n",
    "            bbox_center_x_norm = bbox_center_x / img_width\n",
    "            bbox_center_y_norm = bbox_center_y / img_height\n",
    "\n",
    "            # 框归一化宽度\n",
    "            bbox_width_norm = bbox_width / img_width\n",
    "            # 框归一化高度\n",
    "            bbox_height_norm = bbox_height / img_height\n",
    "\n",
    "            # 生成 YOLO 格式的一行标注，指定保留小数点后几位\n",
    "            bbox_yolo_str = '{} {:.4f} {:.4f} {:.4f} {:.4f}'.format(bbox_class_id, bbox_center_x_norm, bbox_center_y_norm, bbox_width_norm, bbox_height_norm)\n",
    "            # 写入 txt 文件中\n",
    "            f.write(bbox_yolo_str + '\\n')\n",
    "            \n",
    "print('{} --> {} 转换完成'.format(labelme_path, yolo_txt_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d809822-9515-4035-bbb7-33c9bc0fdea2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
